package com.gitee.empty_null.authorization.provider;

import com.gitee.empty_null.authorization.authentication.OAuth2UsernamePasswordAuthenticationToken;
import org.springframework.security.authentication.*;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.security.core.userdetails.UserDetailsChecker;
import org.springframework.security.core.userdetails.UserDetailsService;
import org.springframework.security.core.userdetails.UsernameNotFoundException;
import org.springframework.security.oauth2.core.*;
import org.springframework.security.oauth2.core.endpoint.OAuth2ParameterNames;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;
import org.springframework.security.oauth2.core.oidc.endpoint.OidcParameterNames;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AccessTokenAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContextHolder;
import org.springframework.security.oauth2.server.authorization.token.DefaultOAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;

import java.security.Principal;
import java.util.*;

import static com.gitee.empty_null.authorization.utils.OAuth2AuthenticationProviderUtils.getAuthenticatedClientElseThrowInvalidClient;

/**
 * @author xuhainan
 * @date 2024/1/15 16:40
 * @region hefei
 */
public class OAuth2UsernamePasswordAuthenticationProvider extends AbstractAuthenticationProvider {

    @SuppressWarnings("deprecation")
    private static final AuthorizationGrantType PASSWORD_GRANT_TYPE = AuthorizationGrantType.PASSWORD;
    private final UserDetailsChecker preAuthenticationChecks = new DefaultPreAuthenticationChecks();
    private final UserDetailsChecker postAuthenticationChecks = new DefaultPostAuthenticationChecks();
    private final OAuth2AuthorizationService authorizationService;
    private final UserDetailsService userDetailsService;

    public OAuth2UsernamePasswordAuthenticationProvider(OAuth2AuthorizationService authorizationService,
                                                        OAuth2TokenGenerator<? extends OAuth2Token> tokenGenerator,
                                                        UserDetailsService userDetailsService) {
        super(tokenGenerator);
        Assert.notNull(authorizationService, "authorizationService cannot be null");
        Assert.notNull(tokenGenerator, "tokenGenerator cannot be null");
        this.authorizationService = authorizationService;
        this.userDetailsService = userDetailsService;

    }


    @Override
    public Authentication authenticate(Authentication authentication) throws AuthenticationException {

        OAuth2UsernamePasswordAuthenticationToken resourceOwnerPasswordAuthentication
                = (OAuth2UsernamePasswordAuthenticationToken) authentication;

        OAuth2ClientAuthenticationToken clientPrincipal =
                getAuthenticatedClientElseThrowInvalidClient(resourceOwnerPasswordAuthentication);

        RegisteredClient registeredClient = Optional.ofNullable(clientPrincipal.getRegisteredClient())
                .orElseThrow();

        if (!registeredClient.getAuthorizationGrantTypes().contains(PASSWORD_GRANT_TYPE)) {
            throw new OAuth2AuthenticationException(OAuth2ErrorCodes.UNAUTHORIZED_CLIENT);
        }

        Authentication principal = getUsernamePasswordAuthentication(resourceOwnerPasswordAuthentication);

        // 校验client中scope信息
        Set<String> scopes = validateScopes(resourceOwnerPasswordAuthentication.getScopes(), registeredClient);

        OAuth2Authorization.Builder authorizationBuilder = OAuth2Authorization.withRegisteredClient(registeredClient)
                .principalName(principal.getName())
                .authorizationGrantType(PASSWORD_GRANT_TYPE)
                .authorizedScopes(scopes)
                .attribute(Principal.class.getName(), principal);

        // @formatter:off
        DefaultOAuth2TokenContext.Builder tokenContextBuilder = DefaultOAuth2TokenContext.builder()
                .registeredClient(registeredClient)
                .principal(principal)
                .authorizationServerContext(AuthorizationServerContextHolder.getContext())
                .authorizedScopes(scopes)
                .tokenType(OAuth2TokenType.ACCESS_TOKEN)
                .authorizationGrantType(PASSWORD_GRANT_TYPE)
                .authorizationGrant(resourceOwnerPasswordAuthentication);
        // @formatter:on

        // ----- Access token -----
        OAuth2AccessToken accessToken = generateOAuth2AccessToken(tokenContextBuilder, authorizationBuilder);
        // ----- Refresh token -----
        OAuth2RefreshToken refreshToken = generateOAuth2RefreshToken(
                tokenContextBuilder, authorizationBuilder, registeredClient, clientPrincipal);
        // ----- ID token -----
        OidcIdToken idToken = generateOidcToken(tokenContextBuilder, authorizationBuilder, scopes, principal);

        OAuth2Authorization authorization = authorizationBuilder.build();
        this.authorizationService.save(authorization);

        Map<String, Object> additionalParameters = idTokenAdditionalParameters(idToken);

        return new OAuth2AccessTokenAuthenticationToken(
                registeredClient, clientPrincipal, accessToken, refreshToken, additionalParameters);
    }


    private Map<String, Object> idTokenAdditionalParameters(OidcIdToken idToken) {
        Map<String, Object> additionalParameters = Collections.emptyMap();
        if (idToken != null) {
            additionalParameters = new HashMap<>();
            additionalParameters.put(OidcParameterNames.ID_TOKEN, idToken.getTokenValue());
        }
        return additionalParameters;
    }

    protected Set<String> validateScopes(Set<String> requestedScopes, RegisteredClient registeredClient) {
        Set<String> authorizedScopes = Collections.emptySet();
        if (!CollectionUtils.isEmpty(requestedScopes)) {
            for (String requestedScope : requestedScopes) {
                if (!registeredClient.getScopes().contains(requestedScope)) {
                    throw new OAuth2AuthenticationException(OAuth2ErrorCodes.INVALID_SCOPE);
                }
            }
            authorizedScopes = new LinkedHashSet<>(requestedScopes);
        }
        return authorizedScopes;
    }

    private Authentication getUsernamePasswordAuthentication(OAuth2UsernamePasswordAuthenticationToken authentication) {
        Map<String, Object> additionalParameters = authentication.getAdditionalParameters();
        String username = parseParameter(additionalParameters.get(OAuth2ParameterNames.USERNAME));
        return retrieveUser(username)
                .map(user -> {
                    this.preAuthenticationChecks.check(user);
                    this.additionalAuthenticationChecks(user, additionalParameters);
                    this.postAuthenticationChecks.check(user);
                    return this.createUsernamePasswordAuthenticationToken(user, authentication);
                }).orElseThrow(() -> new UsernameNotFoundException("[" + username + "] not Found"));
    }

    private String parseParameter(Object object) {
        if (object instanceof String str) {
            return str;
        }
        if (object instanceof Collection<?> collection && (!CollectionUtils.isEmpty(collection))) {
            if (collection instanceof List<?> list) {
                return (String) list.get(0);
            }
            if (collection instanceof Set<?> set) {
                return (String) set.iterator().next();
            }
        }
        throw new ClassCastException(object.getClass() + "类型转换失败");
    }

    @Override
    protected Optional<UserDetails> retrieveUser(String username) {
        return Optional.of(userDetailsService.loadUserByUsername(username));
    }

    protected void additionalAuthenticationChecks(UserDetails userDetails, Map<String, Object> additionalParameters)
            throws AuthenticationException {
        String presentedPassword = parseParameter(additionalParameters.get(OAuth2ParameterNames.PASSWORD));
        if (!this.getPasswordEncoder().matches(presentedPassword, userDetails.getPassword())) {
            this.logger.debug("Failed to authenticate since password does not match stored value");
            throw new BadCredentialsException("Invalid credentials");
        }
    }

    private UsernamePasswordAuthenticationToken createUsernamePasswordAuthenticationToken(UserDetails userDetails,
                                                                                          Authentication authentication) {
        UsernamePasswordAuthenticationToken usernamePasswordAuthenticationToken =
                UsernamePasswordAuthenticationToken.authenticated(userDetails, userDetails.getPassword(),
                        userDetails.getAuthorities());
        usernamePasswordAuthenticationToken.setDetails(authentication.getDetails());
        return usernamePasswordAuthenticationToken;
    }


    @Override
    public boolean supports(Class<?> authentication) {
        return OAuth2UsernamePasswordAuthenticationToken.class.isAssignableFrom(authentication);
    }

    private class DefaultPreAuthenticationChecks implements UserDetailsChecker {

        @Override
        public void check(UserDetails user) {
            if (!user.isAccountNonLocked()) {
                OAuth2UsernamePasswordAuthenticationProvider.this.logger
                        .debug("Failed to authenticate since user account is locked");
                throw new LockedException(OAuth2UsernamePasswordAuthenticationProvider.this.messages
                        .getMessage("AbstractUserDetailsAuthenticationProvider.locked",
                                "User account is locked"));
            }
            if (!user.isEnabled()) {
                OAuth2UsernamePasswordAuthenticationProvider.this.logger
                        .debug("Failed to authenticate since user account is disabled");
                throw new DisabledException(OAuth2UsernamePasswordAuthenticationProvider.this.messages
                        .getMessage("AbstractUserDetailsAuthenticationProvider.disabled",
                                "User is disabled"));
            }
            if (!user.isAccountNonExpired()) {
                OAuth2UsernamePasswordAuthenticationProvider.this.logger
                        .debug("Failed to authenticate since user account has expired");
                throw new AccountExpiredException(OAuth2UsernamePasswordAuthenticationProvider.this.messages
                        .getMessage("AbstractUserDetailsAuthenticationProvider.expired",
                                "User account has expired"));
            }
        }

    }

    private class DefaultPostAuthenticationChecks implements UserDetailsChecker {

        @Override
        public void check(UserDetails user) {
            if (!user.isCredentialsNonExpired()) {
                OAuth2UsernamePasswordAuthenticationProvider.this.logger
                        .debug("Failed to authenticate since user account credentials have expired");
                throw new CredentialsExpiredException(OAuth2UsernamePasswordAuthenticationProvider.this.messages
                        .getMessage("AbstractUserDetailsAuthenticationProvider.credentialsExpired",
                                "User credentials have expired"));
            }
        }

    }
}
